﻿#include "WhiteList.h"


PHASHTABLE g_proc_white_list = NULL;


NTSTATUS add_whitelist_rule(PHASHTABLE pTable, PVOID pBuf, ULONG size)
{

	ELEMENT elem = { 0 };
	PHIPS_RULE_NODE pHipsNode = (PHIPS_RULE_NODE)pBuf;


	if (pBuf == NULL || size != sizeof(HIPS_RULE_NODE))
	{
		return FALSE;
	}

	elem.pData = (PVOID)pHipsNode->des_path;
	elem.ulSize = HIPS_PATH_LEN;
	return Insert(pHipsNode->key, &elem, pTable);
}


BOOLEAN Is_In_White_List(PHASHTABLE pHashTable, WCHAR *wszNtPathName)
{
	BOOLEAN bRet = FALSE;
	PTWOWAY pNode = NULL;
	PLIST_ENTRY pListHead = NULL;
	PLIST_ENTRY pListLink = NULL;
	unsigned int i = 0;
	WCHAR *pTmp = NULL;
	UNICODE_STRING usNtPathName = { 0 };
	UNICODE_STRING usResource = { 0 };
	if (NULL == pHashTable || NULL == wszNtPathName || *wszNtPathName == 0)
	{
		return FALSE;
	}

	RtlInitUnicodeString(&usNtPathName, wszNtPathName);

	

	LockRead(pHashTable);
	if (pHashTable != NULL && pHashTable->pListHeads != NULL)
	{
		for (i = 0; i < pHashTable->tableSize; i++)
		{
			pListHead = pListLink = &pHashTable->pListHeads[i];
			if (pListHead == NULL)
			{
				KdPrint(("pListHead is NULL!\n"));
				continue;
			}
			if (!IsListEmpty(pListHead))
			{
				pListLink = pListHead->Flink;
				do
				{
					UNICODE_STRING uni = { 0 };
					WCHAR szBuffer[MYMAXLEN] = { 0 };

					pNode = CONTAINING_RECORD(pListLink, TWOWAY, linkfield);
					pListLink = pListLink->Flink;
					pTmp = (WCHAR *)(pNode->data.pData);
					RtlInitUnicodeString(&uni, pTmp);

					//获取进程版本信息
					//CompanyName 公司名称
					//FileDescription 文件描述
					//InternalName 内部名称
					//ProductName产品名称
					if (!get_version_info_by_path_name(&usNtPathName, L"\\ProductName", szBuffer))
					{
						//获取失败就不在白名单里了
						KdPrint(("get_version_info_by_path_name failed\n"));
						bRet = FALSE;
						goto FUN_END;
					}

					RtlInitUnicodeString(&usResource, szBuffer);
					KdPrint(("white szbuffer:%wZ", &usResource));

					if (FindSubUnicodeStrSafe(&uni, &usResource) != NULL)
					{
						//匹配到了白名单
						bRet = TRUE;
						goto FUN_END;
					}
				} while (pListLink != pListHead);
			}
		}
	}
FUN_END:
	UnLock(pHashTable);
	return bRet;
}


void InitWhiteList()
{
	//默认的白名单
	
	PHIPS_RULE_NODE pWhiteBuf = NULL;
	//手动添加默认白名单
	//Microsoft® Windows® Operating System->CRC32：a163c5db
	pWhiteBuf = (PHIPS_RULE_NODE)ExAllocatePoolWithTag(PagedPool, sizeof(HIPS_RULE_NODE), 'blck');
	RtlZeroMemory(pWhiteBuf, sizeof(HIPS_RULE_NODE));
	RtlCopyMemory(pWhiteBuf->des_path, L"Microsoft® Windows® Operating System", sizeof(L"Microsoft® Windows® Operating System") - 2);
	pWhiteBuf->key = 0xa163c5db;
	add_whitelist_rule(g_proc_white_list, (PVOID)pWhiteBuf, sizeof(HIPS_RULE_NODE));
}



BOOLEAN is_trust_process(PHIPS_RULE_NODE fi)
{

	//C : \Program Files\VMware\VMware Tools\vmtoolsd.exe
	//RtlCompareUnicodeString C:\Program Files\VMware\VMware Tools\VMwareTray.exe

	/*
	C:\Windows\System32\svchost
	C:\Windows\System32\grpconv
	C:\Windows\explorer
	C:\Windows\conhost.exe
	*/
	
	//加入自己的R3进程名..
	WCHAR *wszTrustName[] = { R3NAME, L"vmtoolsd.exe", L"VMwareTray.exe", L"SearchIndexer.exe", L"SearchFilterHost.exe", 
		L"SearchProtocolHost.exe",
		L"explorer.exe",
		L"svchost.exe"};
	WCHAR *wszTrustPath[] = { L"C:\\Windows\\System32\\svchost.exe", L"C:\\Windows\\System32\\grpconv.exe", L"C:\\Windows\\explorer.exe", L"C:\\Windows\\System32\\conhost.exe", L"C:\\Windows\\System32\\smss.exe", L"C:\\Windows\\System32\\services.exe" ,L"C:\\Windows\\System32\\csrss.exe"};
	
	WCHAR *wszTrustDest[] = { L"C:\\Windows\\System32\\imm32.dll", L"C:\\Windows\\System32\\apphelp.dll", L"C:\\Windows\\System32\\winbrand.dll" };

	PEPROCESS pSelfEprocess = NULL;
	UNICODE_STRING usSrcName = { 0 };
	UNICODE_STRING usDesName = { 0 };
	UNICODE_STRING usSvchost = { 0 };
	ULONG i = 0;
	ULONG uCountName = sizeof(wszTrustName) / sizeof(ULONG);
	ULONG uCountPath = sizeof(wszTrustPath) / sizeof(ULONG);
	ULONG uCountTrustDest = sizeof(wszTrustDest) / sizeof(ULONG);

	NTSTATUS st = 0;
	UCHAR *uSelfName = NULL;
	BOOLEAN bRet = FALSE;

	if (fi->des_path == NULL)
		return FALSE;
	if (*fi->des_path == 0)
		return FALSE;

	RtlInitUnicodeString(&usSrcName, fi->src_path);
	RtlInitUnicodeString(&usDesName, fi->des_path);

	

	
	
	
	for (i = 0; i < uCountName; ++i)
	{
		RtlInitUnicodeString(&usSvchost, wszTrustName[i]);
		if (FindSubUnicodeStrSafe(&usSrcName, &usSvchost) != NULL)
		{
			KdPrint(("pass RtlCompareUnicodeString %wZ", &usSrcName));
			bRet = TRUE;
			goto FUN_END;
		}
	}
	


	//svchost.exe放过
	for (i = 0; i < uCountPath; ++i)
	{

		RtlInitUnicodeString(&usSvchost, wszTrustPath[i]);

		if (0 == RtlCompareUnicodeString(&usSvchost, &usSrcName, FALSE))
		{
			KdPrint(("pass RtlCompareUnicodeString %wZ", &usSrcName));
			bRet = TRUE;
			goto FUN_END;
		}
	}

	
	for (i = 0; i < uCountTrustDest; ++i)
	{
		RtlInitUnicodeString(&usSvchost, wszTrustDest[i]);
		if (RtlCompareUnicodeString(&usDesName, &usSvchost,FALSE) == 0)
		{
			KdPrint(("pass RtlCompareUnicodeString %wZ", &usSrcName));
			bRet = TRUE;
			goto FUN_END;
		}
	}
	


FUN_END:
	return bRet;
}